import os
import torch
from vlmeval.config import supported_VLM

# 手动指定 CUDA 路径
os.environ["CUDA_HOME"] = "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v12.9"
os.environ["PATH"] = os.environ["CUDA_HOME"] + "\\bin;" + os.environ["PATH"]

# 初始化 CUDA
if not torch.cuda.is_available():
    torch.cuda.init()

print(f"CUDA 可用: {torch.cuda.is_available()}")
print(f"CUDA 版本: {torch.version.cuda}")

model = supported_VLM["llava_v1.5_7b"]()

# # 前向单张图片
# ret = model.generate(['assets/apple.jpg', 'What is in this image?'])
# print(ret)  # 这张图片上有一个带叶子的红苹果
# # 前向多张图片
# ret = model.generate(['assets/apple.jpg', 'assets/apple.jpg', 'How many apples are there in the provided images? '])
# print(ret)  # 提供的图片中有两个苹果

# 测试单张图片
print("\n测试单张图片...")
ret = model.generate(['assets/apple.jpg', '这张图片里有什么?'])
print(f"结果: {ret}")

# 测试多张图片
print("\n测试多张图片...")
ret = model.generate(['assets/apple.jpg', 'assets/apple.jpg', '提供的图片中有多少个苹果?'])
print(f"结果: {ret}")